Skip to content

Return Euclidean distances from KMeans transform - #8539

Open
Hashim1999164 wants to merge 3 commits into
NVIDIA:mainfrom
Hashim1999164:fix/kmeans-transform-euclidean
Open

Return Euclidean distances from KMeans transform#8539
Hashim1999164 wants to merge 3 commits into
NVIDIA:mainfrom
Hashim1999164:fix/kmeans-transform-euclidean

Conversation

@Hashim1999164

@Hashim1999164 Hashim1999164 commented Sep 1, 2026

Copy link
Copy Markdown

Summary

KMeans.transform currently returns squared distances. sklearn returns Euclidean distances.

This change takes the square root of the transform output so the Python API matches sklearn. score is unchanged.

Fixes #8536

Test plan

Added a regression test against the six point example from the issue.

@Hashim1999164
Hashim1999164 requested a review from a team as a code owner September 1, 2026 16:14
@copy-pr-bot

copy-pr-bot Bot commented Sep 1, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@github-actions github-actions Bot added the Cython / Python Cython or Python issue label Sep 1, 2026
@coderabbitai

coderabbitai Bot commented Sep 1, 2026

Copy link
Copy Markdown

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 5f461c3e-3729-4f20-90f8-09238b03443c

📥 Commits

Reviewing files that changed from the base of the PR and between f971d8f and ad513a8.

📒 Files selected for processing (2)
  • python/cuml/cuml/cluster/kmeans.pyx
  • python/cuml/tests/test_kmeans.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • python/cuml/tests/test_kmeans.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.


📝 Summary

Summary by CodeRabbit

  • Bug Fixes

    • KMeans transformations now return Euclidean distances rather than squared-Euclidean distances.
    • Improved consistency with scikit-learn distance results for transformed data.
  • Tests

    • Added regression coverage validating KMeans distance calculations against expected values and scikit-learn results.

Walkthrough

KMeans.transform now requests Euclidean distances from the backend. A regression test validates expected values and scikit-learn compatibility.

Changes

KMeans distance correction

Layer / File(s) Summary
Backend distance selection and regression coverage
python/cuml/cuml/cluster/kmeans.pyx, python/cuml/tests/test_kmeans.py
KMeans.transform uses DistanceType.L2SqrtExpanded and returns the backend output unchanged. The regression test compares fixed Euclidean distances with explicit values and scikit-learn.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to ad513

KMeans.transform now returns Euclidean distances compatible with scikit-learn, with regression coverage for the corrected output. No current merge-blocking risk remains.

Suggested reviewers: jcrist

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 1 files. (1 skipped: 1… Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly and concisely states the primary change: KMeans.transform now returns Euclidean distances.
Description check ✅ Passed The description accurately explains the squared-distance bug, the scikit-learn compatibility goal, the unchanged score behavior, and the regression test.
Linked Issues check ✅ Passed The implementation requests square-rooted Euclidean distances from the backend, and the regression test verifies expected values against explicit calculations and scikit-learn. This satisfies the obje…
Out of Scope Changes check ✅ Passed The changes are limited to the KMeans.transform implementation and its regression test. Both changes directly support the linked issue objective.
Full details: Docstring Coverage

Explanation

Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 1 files. (1 skipped: 1 unsupported.)

  • Fix all pre-merge checks with AI
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@viclafargue viclafargue left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for opening a PR!

Comment thread python/cuml/cuml/cluster/kmeans.pyx Outdated
@viclafargue viclafargue added bug Something isn't working non-breaking Non-breaking change labels Sep 4, 2026

@viclafargue viclafargue left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! LGTM

@viclafargue

Copy link
Copy Markdown
Contributor

/ok to test ad513a8

@viclafargue

Copy link
Copy Markdown
Contributor

@Hashim1999164, the sklearn.cluster.tests.test_k_means::test_transform[42-KMeans] test from the Scikit-Learn test suite is now passing with cuml.accel. Could you remove it from the xfail list over there?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working Cython / Python Cython or Python issue non-breaking Non-breaking change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] KMeans.transform returns squared distances instead of Euclidean distances

3 participants